'''
* This is the projet for Brtc LlmOps Platform
* @Author Leon-liao <liaosiliang@alltman.com>
* @Description //TODO 
* @File: 4_study_prompt_route.py
* @Time: 2025/10/31
* @All Rights Reserve By Brtc
'''
import dotenv
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables import RunnablePassthrough, RunnableLambda
from langchain_openai import OpenAIEmbeddings, ChatOpenAI

dotenv.load_dotenv()
# 1.定义两份不同的prompt模板(物理模板、数学模板)
physics_template = """你是一位非常聪明的物理教程。
你擅长以简洁易懂的方式回答物理问题。
当你不知道问题的答案时，你会坦率承认自己不知道。

这是一个问题：
{query}"""
math_template = """你是一位非常优秀的数学家。你擅长回答数学问题。
你之所以如此优秀，是因为你能将复杂的问题分解成多个小步骤。
并且回答这些小步骤，然后将它们整合在一起回来更广泛的问题。

这是一个问题：
{query}"""

#2、创建文本嵌入模型
embeddings = OpenAIEmbeddings(model="text-embedding-3-small")
llm = ChatOpenAI(model="gpt-4o-mini")
prompt_templates = [math_template,physics_template]

def prompt_router(input)->ChatPromptTemplate:
    """根据传递的query计算返回不同的提示词模板"""
    #1、计算传入的query 返回不同的提示词模板
    query_prompt = ChatPromptTemplate.from_messages([
        ("system",""" 
            #角色:
            你是一个问题分类专家，你需要将用户的问题分类成数学问题或者是物理问题，如果用的问题是数学问题请回复0， 如果用户的问问题是物理问日请回复1\n
            #示例:
            Human:请问什么是牛顿第一定律
            AI:1
            Human:请问什么是一元二次方程
            AI:0
            #限制:
            请严格按照示例回复用户的问题， 不要废话，直接输出答案！"""),
        ("human","{query}")
    ])
    query_chain = query_prompt|llm|StrOutputParser()
    res = query_chain.invoke({"query":input["query"]})
    if res=="1":
        print("使用物理模板！")
    else:
        print("使用数学模板！")

    return ChatPromptTemplate.from_template(prompt_templates[int(res)])


chain = (
    {"query":RunnablePassthrough()}
    |RunnableLambda(prompt_router)
    |llm
    |StrOutputParser()
)

print(chain.invoke("什么是黑洞？"))
print("==============================================")
print(chain.invoke("什么是一元二次方程"))